import os
import asyncio
import openpyxl
from io import BytesIO
from threading import Lock
from datetime import datetime
from PIL import Image as PILImage
from config.settings import sizeChange
from playwright.async_api import async_playwright

# 初始化锁
lock = Lock()
browser_instance = None

# 文件名合法化
def sanitize_filename(filename):
    filename = filename.replace('.', '-')
    return "".join(c for c in filename if c.isalnum() or c in (' ', '-', '_'))

# 生产者：获取页面截图和标题，并放入队列
# 同步截图函数
async def screen_shoot(url, queue, resize_ratio=sizeChange):

    try:
        # 使用 async_playwright 来启动 Playwright
        async with async_playwright() as p:
            browser = await p.chromium.launch()  # 异步启动浏览器
            context = await browser.new_context()  # 异步创建上下文

            try:
                page = await context.new_page()  # 异步创建页面
                # 跳过不重要的资源请求
                await page.route('**/*', lambda route: route.abort() if route.request.resource_type in ['font', 'image'] else route.continue_())

                # 访问 URL，确保页面加载完成
                await page.goto(url, wait_until='commit')
                await page.wait_for_load_state('domcontentloaded')
                await page.wait_for_load_state('networkidle')

                # 获取网页标题
                page_title = await page.title() if await page.title() else '无标题'

                # 截图并保存为图片
                screenshot = await page.screenshot()
                image = PILImage.open(BytesIO(screenshot))  # 使用 PIL.Image 的 open 方法

                # 根据 resize_ratio 调整图片大小
                if resize_ratio != 1.0:
                    original_width, original_height = image.size
                    new_width = int(original_width * resize_ratio)
                    new_height = int(original_height * resize_ratio)
                    image = image.resize((new_width, new_height))
                # 将获取的数据放入队列
                queue.put((url, page_title, image))
            except Exception as page_error:
                print(f"创建页面或处理页面时出错: {page_error}")

            finally:
                # 确保关闭页面和浏览器
                try:
                    if 'page' in locals():
                        await page.close()
                    if browser:
                        await browser.close()
                except Exception as close_error:
                    print(f"关闭页面或浏览器时出错: {close_error}")

    except Exception as e:
        print(f"Error accessing {url}: {e}")

# 消费者：获取队列中的数据并将其写入Excel文件
def data_write(queue):
    data_buffer = []
    while True:
        # 从队列中取出数据
        data = queue.get()
        if data is None:
            break  # 如果队列为空且收到None信号，退出线程

        # 将数据缓存
        data_buffer.append(data)

        # 每处理完一条数据就标记任务完成
        queue.task_done()

        # 如果缓存满了或者队列有数据了，开始写入文件
        if len(data_buffer) > 0:
            write_to_excel(data_buffer)
            data_buffer.clear()  # 清空缓存

def write_to_excel(data_buffer):
    with lock:
        # 获取当前日期作为文件夹名称
        current_date = datetime.now().strftime("%m%d")  # 月日格式
        folder_path = os.path.join(os.getcwd(), current_date)

        # 如果文件夹不存在则创建
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)

        # 设置Excel文件路径
        excel_file = os.path.join(folder_path, f"{sanitize_filename(current_date)}_data.xlsx")

        # 如果Excel文件不存在，创建新的
        if not os.path.exists(excel_file):
            workbook = openpyxl.Workbook()
            sheet = workbook.active
            sheet.title = "Web Data"

            # 添加表头
            sheet.append(["URL", "Title", "Screenshot"])

            # 保存初始的Excel文件
            workbook.save(excel_file)

        # 打开Excel文件进行写入
        workbook = openpyxl.load_workbook(excel_file)
        sheet = workbook.active

        # 将数据写入Excel文件
        for url, title, screenshot in data_buffer:
            # 生成合法的图片文件名：url-标题
            file_name = f"{sanitize_filename(url)}-{sanitize_filename(title[:10])}.png"  # 取标题前10个字符
            screenshot_file_path = os.path.join(folder_path, file_name)  # 文件路径

            # 保存截图为PNG文件
            screenshot.save(screenshot_file_path)

            # 添加数据行，URL，标题和截图的文件路径
            sheet.append([url, title, screenshot_file_path])

        # 保存Excel文件
        workbook.save(excel_file)

async def operation_on_success(url, queue):
    await screen_shoot(url, queue)
